#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
C. jejuni Penner（HS）in-silico 分型 — v2_tidy（纵向输出）
改进点：
- 更严格阈值：错配 ≤1（MAX_MISMATCH=1），片段大小容忍 ±5 bp（SIZE_TOL=5）
- 去重：同一 sample+marker 仅保留与理论产物最接近的一条
- contig 聚类：优先选择携带最多 marker 的 contig（近似 CPS 位点）
- 直接输出纵向（tidy）三件套：
    1) hs_typing_summary.tsv   （每样本一行：primary_HS、all_HS、n_HS、lpxA_present）
    2) hs_calls_long.tsv       （每样本×HS 一行：hs_rank 表示主次顺序）
    3) amplicons_tidy.tsv      （去重+聚类后的条带明细）

用法：
1) 把 .fa/.fasta/.fna 放到 ASSEMBLIES_DIR 目录（见下方常量）
2) 运行：python hs_typing_insilico_standalone_v2_tidy.py
"""

# ===== 固定参数（可直接在此处修改）=====
ASSEMBLIES_DIR = "./assemblies"
MAX_MISMATCH   = 1
SIZE_TOL       = 5
FA_EXTS        = (".fa", ".fasta", ".fna")

OUT_SUMMARY    = "hs_typing_summary.tsv"
OUT_HS_LONG    = "hs_calls_long.tsv"
OUT_AMP_TIDY   = "amplicons_tidy.tsv"

# ========= 引物表（Mix, Marker, 预期bp, HS列表, Forward, Reverse）=========
PRIMERS = [
    # Mix Alpha
    ("Alpha","Mu_HS2",     62,  ["HS2"],                  "CAGCATTGGAGGATTTACAATATAT",               "CATCCTAGCACAACTCACTTCA"),
    ("Alpha","Mu_HS3",    149,  ["HS3"],                  "GGTAAGGTTGATTCTGGGTTTAAT",                 "AGATTAGGCCAAGCAATGATAA"),
    ("Alpha","Mu_HS4A",   370,  ["HS4A"],                 "TATATTTGGTTAGGGATCCA",                     "CCTAACATATCATACACTACGGT"),
    ("Alpha","Mu_HS6",    185,  ["HS6","HS7"],            "CATACATTTG CTTTCAGATTCTTTAC".replace("\u2003",""), "ACACGCCTATTGTTGTTGTTC"),
    ("Alpha","Mu_HS10",   229,  ["HS10"],                 "TCTTATGCAGCACGCTGAT",                       "CAAATTCAATCGACTAGCCACT"),
    ("Alpha","Mu_HS15",   325,  ["HS15","HS31","HS58"],   "ACAGGTAATAAAATGTGCGAGTTT",                 "ATGCATCTGCAACATCATCC"),
    ("Alpha","Mu_HS41",   279,  ["HS41"],                 "CTTACATATGCTGGTAGAGATGATATG",               "TGCAATCTCTAAAGCCCAAG"),
    ("Alpha","Mu_HS53",   251,  ["HS53"],                 "AGGCAAGCAGGAATTGTTT",                       "TTAATTGCTCTTTGGCAATCTT"),
    ("Alpha","Mu_HS19",   450,  ["HS19"],                 "CGAGGATGAAAATGCCTCAA",                      "GGCAACAAACAAACATATTCAGA"),
    ("Alpha","Mu_HS63",   522,  ["HS63"],                 "AAATTTGTTTTTCATATTTTTACGG",                 "TTAGGTGCGGTTACCAAAGG"),
    ("Alpha","Mu_HS33",   819,  ["HS33","HS35"],          "GTAGCGGATCAGCAGCATTA",                      "CATCAAAATCATCTTTTAACACCAA"),
    # Mix Beta
    ("Beta","Mu_HS1",     610,  ["HS1","HS44"],           "TTGGCGGTAAGTTTTTGAAGA",                     "GCAAGAGAAACATCTCGCCTA"),
    ("Beta","Mu_HS4B",    652,  ["HS4B"],                 "GTGGACATGGAACTGGGACT",                       "AAAACGTTTAAAGTCAGTGGAAA"),
    ("Beta","Mu_HS8",     342,  ["HS8","HS17","HS32"],    "TTCACGTGGAGGATTATTGG",                      "TTGAACATTTCATGTGTATTCCCTA"),
    ("Beta","Mu_HS23/36", 161,  ["HS23","HS36"],          "GCTTGGGAGATGAATTTACCTTTA",                   "GCTTTATATCTATCCAGTCCATTATCA"),
    ("Beta","Mu_HS42",    440,  ["HS42"],                 "ATGGTAAAACCGGCATTTCA",                       "ATGCTTCAGTTCCACCCAAA"),
    ("Beta","Mu_HS57",    100,  ["HS57"],                 "GGGGTAAAATAGCCAATATTCCA",                    "CCAACAAGCCATATTTGTTTTTC"),
    ("Beta","Mu_HS12",    201,  ["HS12"],                 "GGAGGTAAAACGATATTCTCCTTAAA",                 "TGAAGATTTTGAATGGATGTGTG"),
    ("Beta","Mu_HS27",    280,  ["HS27"],                 "GAATAAATATTGCTTCCATACTTTCAA",               "GCAAAATGAGAATCTCCACCA"),
    ("Beta","Mu_HS21",    801,  ["HS21"],                 "TGGATGGGATATTGATGACAA",                      "CCCTGGAAGAGTATGGGACA"),
    ("Beta","Mu_HS31",    857,  ["HS5","HS31"],           "GGCAAAGAGCTTTATTTTGTTGA",                    "GCCGTAGCAACATCAAATACA"),
    # Mix Gamma
    ("Gamma","Mu_HS44",   148,  ["HS44"],                 "AGAAGATGCACTAGGCTCTAG",                      "GCTATCTAATTCCATCCCTG"),
    ("Gamma","Mu_HS45",   128,  ["HS45","HS5","HS32","HS60"], "TCCACTTGGGATGAAAAGGA",                 "ACCGCATACTTTGAGCCTGT"),
    ("Gamma","Mu_HS29",   185,  ["HS29"],                 "CCCATATTTAAACAATGGAGTGA",                    "TCATACTTTGAAAAACATTATCTGGA"),
    ("Gamma","Mu_HS22",   216,  ["HS22"],                 "TCATGGAGCTGGAACAACAG",                        "GCTGGAACTTCTTTTGCAATC"),
    ("Gamma","Mu_HS9",    278,  ["HS9"],                  "AAAACTATTAGCTTGATTTTACCTTGG",                "GCGAAAGACGGATTGTTCAT"),
    ("Gamma","Mu_HS37",   541,  ["HS37"],                 "TGGATGAAGGGGACTTATGG",                        "TGG TTTGAAGAGCATCAGCA".replace("\u2003","")),
    ("Gamma","Mu_HS18",   653,  ["HS18"],                 "CAGCTATAAATCATGGGTATTGGA",                    "GTAATCAATACATTTTTCCTTGCTT"),
    # lpxA 内控（有简并碱基）
    ("Gamma","lpxA",      331,  ["C.jejuni"],             "ACAACTTGGTGACGATGTTGTA",                      "CAATCATGDGCDATATGASAATAHGCCAT"),
    # Mix Delta
    ("Delta","Mu_HS58",    85,  ["HS58","HS32"],          "TCCGGAAAAATTTTATTTAGATTCTC",                  "AACAATACCAGGATACCAATCTTCA"),
    ("Delta","Mu_HS52",   170,  ["HS52"],                 "AAAACACGCTATTAATCATGGTGAC",                   "ATGTAGGCCAAGTTATACAACCTTTT"),
    ("Delta","Mu_HS60",   241,  ["HS60"],                 "GAAATCATTTTTATGATATTGTGGTT",                  "TCACAGTCACAATAAATAGCCAAA"),
    ("Delta","Mu_HS55",   341,  ["HS55"],                 "GAGATGGTGGTGGTCATCAA",                         "ACG TTGCAACCAATCCTTTG".replace("\u2003","")),
    ("Delta","Mu_HS32",   420,  ["HS32"],                 "GCATACCAGATGGCTTTGG",                          "AATGCAGCGTGCTTCTTATTT"),
    ("Delta","Mu_HS11",   540,  ["HS11"],                 "GAATTGGACATAACCACGGAAT",                       "ATGCAAAGTGCACATATT CTCC".replace("\u2003","")),
    ("Delta","Mu_HS40",   636,  ["HS40"],                 "CAACCCTTGGATGACAATAGAGA",                      "ACCGTCAATATCATCAGGATTTA"),
    ("Delta","Mu_HS38",   741,  ["HS38"],                 "GCCGCAGGAGATAATGAAGA",                          "TTTGCCTTTTAGATCTTGAGGA"),
]

from typing import List, Dict, Tuple
import os, sys

IUPAC = {
    "A":{"A"}, "C":{"C"}, "G":{"G"}, "T":{"T"}, "U":{"T"},
    "R":{"A","G"}, "Y":{"C","T"}, "S":{"G","C"}, "W":{"A","T"},
    "K":{"G","T"}, "M":{"A","C"}, "B":{"C","G","T"}, "D":{"A","G","T"},
    "H":{"A","C","T"}, "V":{"A","C","G"}, "N":{"A","C","G","T"}
}
def _match_base(gb: str, pb: str) -> bool:
    gb = gb.upper(); pb = pb.upper()
    return gb in IUPAC.get(pb, {pb})

def revcomp(seq: str) -> str:
    comp = str.maketrans("ACGTRYSWKMBDHVNacgtryswkmbdhvn", "TGCAYRSWMKVHDBNtgcayrswmkvhdbn")
    return seq.translate(comp)[::-1]

def read_fasta(path: str):
    name=None; buf=[]
    with open(path, "r", encoding="utf-8", errors="ignore") as fh:
        for line in fh:
            if line.startswith(">"):
                if name is not None:
                    yield name, "".join(buf).replace("\n","").replace("\r","")
                name = line.strip().lstrip(">")
                buf=[]
            else:
                buf.append(line.strip())
        if name is not None:
            yield name, "".join(buf).replace("\n","").replace("\r","")

def find_approx(seq: str, query: str, max_mm: int) -> List[int]:
    s = seq.upper(); q = query.upper(); L = len(q); out=[]
    for i in range(0, len(s)-L+1):
        mm=0
        for a,b in zip(s[i:i+L], q):
            if not _match_base(a,b):
                mm += 1
                if mm > max_mm: break
        if mm <= max_mm: out.append(i)
    return out

def scan_one(fasta_path: str):
    try:
        import pandas as pd
    except ImportError:
        print("[ERROR] 需要 pandas：pip install pandas")
        sys.exit(1)

    rows=[]
    for cname, seq in read_fasta(fasta_path):
        S=seq.upper()
        for mix, name, exp, penner, fwd, rev in PRIMERS:
            R = revcomp(rev)
            f = find_approx(S, fwd, MAX_MISMATCH)
            r = find_approx(S, R,   MAX_MISMATCH)
            for i in f:
                for j in r:
                    if j>i:
                        size = j + len(R) - i
                        if abs(size-exp) <= SIZE_TOL:
                            rows.append([cname,mix,name,exp,size,"+/-"])
            FR = revcomp(fwd)
            r2 = find_approx(S, rev, MAX_MISMATCH)
            f2 = find_approx(S, FR, MAX_MISMATCH)
            for i in r2:
                for j in f2:
                    if j>i:
                        size = j + len(FR) - i
                        if abs(size-exp) <= SIZE_TOL:
                            rows.append([cname,mix,name,exp,size,"-/+"])
    import pandas as pd
    df = pd.DataFrame(rows, columns=["contig","mix","marker","exp_bp","obs_bp","orientation"])
    if df.empty: return df, set()

    df["abs_diff"] = (df["obs_bp"] - df["exp_bp"]).abs()
    df = df.sort_values(["marker","abs_diff","contig"]).drop_duplicates(["marker"], keep="first")

    sub = df[df["marker"]!="lpxA"]
    if not sub.empty:
        top = sub["contig"].value_counts().index[0]
        df = df[(df["contig"]==top) | (df["marker"]=="lpxA")]

    return df, set(df["marker"])

def decide(mset: set) -> List[str]:
    mh = {m: True for m in mset}
    out = []
    if "Mu_HS32" in mh: out.append("HS32")
    elif "Mu_HS8" in mh: out.append("HS8/17")
    if "Mu_HS31" in mh:
        if "Mu_HS45" in mh: out.append("HS5")
        elif "Mu_HS15" in mh: out.append("HS31")
        else: out.append("HS5/31")
    if "Mu_HS1" in mh: out.append("HS44" if "Mu_HS44" in mh else "HS1")
    if "Mu_HS58" in mh: out.append("HS58")
    elif "Mu_HS15" in mh and "Mu_HS31" not in mh: out.append("HS15")
    if "Mu_HS45" in mh and "Mu_HS31" not in mh and "Mu_HS32" not in mh: out.append("HS45")
    if "Mu_HS4A" in mh: out.append("HS4-complex(4A+4B)" if "Mu_HS4B" in mh else "HS4-complex")
    for m,hs in [
        ("Mu_HS2","HS2"),("Mu_HS3","HS3"),("Mu_HS6","HS6/7"),
        ("Mu_HS10","HS10"),("Mu_HS41","HS41"),("Mu_HS53","HS53"),
        ("Mu_HS19","HS19"),("Mu_HS63","HS63"),("Mu_HS33","HS33/35"),
        ("Mu_HS23/36","HS23/36"),("Mu_HS42","HS42"),("Mu_HS57","HS57"),
        ("Mu_HS12","HS12"),("Mu_HS27","HS27"),("Mu_HS21","HS21"),
        ("Mu_HS29","HS29"),("Mu_HS22","HS22"),("Mu_HS9","HS9"),
        ("Mu_HS37","HS37"),("Mu_HS18","HS18"),("Mu_HS52","HS52"),
        ("Mu_HS60","HS60"),("Mu_HS55","HS55"),("Mu_HS11","HS11"),
        ("Mu_HS40","HS40"),("Mu_HS38","HS38")
    ]:
        if m in mh: out.append(hs)
    order=["HS1","HS44","HS2","HS3","HS4-complex","HS4-complex(4A+4B)","HS5","HS31","HS5/31","HS6/7",
           "HS8/17","HS9","HS10","HS11","HS12","HS15","HS18","HS19","HS21","HS22","HS23/36","HS27",
           "HS29","HS32","HS33/35","HS37","HS38","HS40","HS41","HS42","HS52","HS53","HS55","HS57","HS58","HS60","HS63"]
    uniq=[]
    for c in order:
        if c in out and c not in uniq: uniq.append(c)
    for c in out:
        if c not in uniq: uniq.append(c)
    return uniq

def main():
    try:
        import pandas as pd
    except ImportError:
        print("[ERROR] 需要 pandas：pip install pandas")
        sys.exit(1)

    if not os.path.isdir(ASSEMBLIES_DIR):
        print(f"[ERROR] 找不到目录：{ASSEMBLIES_DIR}")
        sys.exit(1)
    files = [fn for fn in os.listdir(ASSEMBLIES_DIR) if fn.lower().endswith(FA_EXTS)]
    if not files:
        print("[WARN] assemblies 目录为空。仍会写出空表。")
        pd.DataFrame(columns=["sample","primary_HS","all_HS","n_HS","lpxA_present"]).to_csv(OUT_SUMMARY, sep="\t", index=False)
        pd.DataFrame(columns=["sample","hs_rank","hs_call","lpxA_present"]).to_csv(OUT_HS_LONG, sep="\t", index=False)
        pd.DataFrame(columns=["sample","contig","mix","marker","exp_bp","obs_bp","orientation"]).to_csv(OUT_AMP_TIDY, sep="\t", index=False)
        return

    all_amp = []
    summary_rows = []
    for fn in sorted(files):
        sample = os.path.splitext(fn)[0]
        df, mset = scan_one(os.path.join(ASSEMBLIES_DIR, fn))
        if not df.empty:
            tmp = df.copy(); tmp.insert(0, "sample", sample); all_amp.append(tmp)
        calls = decide(mset)
        primary = calls[0] if calls else ""
        allc    = ";".join(calls) if calls else ""
        nhs     = len(calls)
        lpx     = "Yes" if "lpxA" in mset else "No"
        summary_rows.append([sample, primary, allc, nhs, lpx])

    amp = pd.concat(all_amp, ignore_index=True) if all_amp else pd.DataFrame(columns=["sample","contig","mix","marker","exp_bp","obs_bp","orientation"])
    amp = amp[["sample","contig","mix","marker","exp_bp","obs_bp","orientation"]].sort_values(["sample","marker","contig","obs_bp"])
    amp.to_csv(OUT_AMP_TIDY, sep="\t", index=False)

    summary = pd.DataFrame(summary_rows, columns=["sample","primary_HS","all_HS","n_HS","lpxA_present"]).sort_values("sample")
    summary.to_csv(OUT_SUMMARY, sep="\t", index=False)

    rows=[]
    for _, r in summary.iterrows():
        sample = r["sample"]; lpx = r["lpxA_present"]
        calls  = [c for c in str(r["all_HS"]).split(";") if c]
        if not calls:
            rows.append([sample, 0, "", lpx])
        else:
            for i, c in enumerate(calls, start=1):
                rows.append([sample, i, c, lpx])
    hs_long = pd.DataFrame(rows, columns=["sample","hs_rank","hs_call","lpxA_present"])
    hs_long.to_csv(OUT_HS_LONG, sep="\t", index=False)

    print("[OK] 写出：", OUT_SUMMARY, OUT_HS_LONG, OUT_AMP_TIDY)
    print(f"[参数] 目录={ASSEMBLIES_DIR}，错配≤{MAX_MISMATCH}，size_tol±{SIZE_TOL} bp")

if __name__ == "__main__":
    main()
